# flow.py
from pocketflow import AsyncFlow
from .nodes import (
    ParsePDFNode, 
    ChunkTextNode, 
    ExtractCheckpointsNode, 
    MergeCheckpointsNode,
    ProofMapperNode, 
    ExportJSONNode,
    SaveJSONNode,
    InteractiveChatNode,
    FinalizeJSONNode,
    LoadExistingJSONNode,
    ValidateJSONNode
)

def create_proofmapper_flow():
    """
    创建ProofMapper异步并行流程，支持可选的交互式修改
    
    Returns:
        AsyncFlow: 配置好的ProofMapper异步流程，支持多线程并行处理和交互式修改
    """
    # 创建核心处理节点
    parse_pdf_node = ParsePDFNode()                    # 同步：解析PDF
    chunk_text_node = ChunkTextNode()                  # 同步：分割文本为135字符小块  
    extract_checkpoints_node = ExtractCheckpointsNode() # 异步并行：提取评估项
    merge_checkpoints_node = MergeCheckpointsNode()    # 异步并行：按模块分组去重评估项
    proofmapper_node = ProofMapperNode()               # 异步并行：生成证明材料清单
    export_json_node = ExportJSONNode()                # 同步：格式化JSON数据
    save_json_node = SaveJSONNode()                    # 同步：保存JSON文件
    
    # 创建交互式修改节点
    interactive_chat_node = InteractiveChatNode()     # 异步：交互式对话修改
    finalize_json_node = FinalizeJSONNode()           # 同步：更新元数据和清理
    final_save_node = SaveJSONNode()                  # 同步：保存最终JSON
    
    # 构建主流程链
    parse_pdf_node >> chunk_text_node >> extract_checkpoints_node >> merge_checkpoints_node >> proofmapper_node >> export_json_node
    
    # 配置保存和交互分支
    export_json_node - "save_and_interactive" >> save_json_node
    export_json_node - "save_and_complete" >> save_json_node
    
    # 保存后的分支
    save_json_node - "interactive" >> interactive_chat_node
    save_json_node - "complete" >> None  # 非交互模式直接结束
    
    # 交互式流程完成后的最终化
    interactive_chat_node - "finalize" >> finalize_json_node
    finalize_json_node - "save" >> final_save_node
    
    # 创建异步流程
    return AsyncFlow(start=parse_pdf_node)


def create_quick_interactive_flow():
    """
    创建快速交互模式流程，直接从现有JSON文件开始
    
    Returns:
        AsyncFlow: 配置好的快速交互流程
    """
    # 创建快速交互模式节点
    load_existing_json_node = LoadExistingJSONNode()  # 同步：加载现有JSON文件
    validate_json_node = ValidateJSONNode()           # 同步：验证JSON格式
    interactive_chat_node = InteractiveChatNode()     # 异步：交互式对话修改
    finalize_json_node = FinalizeJSONNode()           # 同步：更新元数据和清理
    final_save_node = SaveJSONNode()                  # 同步：保存最终JSON
    
    # 构建快速交互流程链
    load_existing_json_node - "validate" >> validate_json_node
    
    # 配置验证分支
    validate_json_node - "direct_interactive" >> interactive_chat_node
    validate_json_node - "invalid" >> None  # 验证失败时结束流程
    
    # 交互式流程完成后的最终化和保存
    interactive_chat_node - "finalize" >> finalize_json_node
    finalize_json_node - "save" >> final_save_node
    
    # 创建异步流程
    return AsyncFlow(start=load_existing_json_node)